use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use numpy::{PyArray, PyArrayDyn, Ix3};
use image::{DynamicImage, GenericImageView, RgbaImage};
use image::imageops;
use anyhow::Result;

/// 灰度处理
fn process_grayscale(img: &DynamicImage) -> RgbaImage {
    img.grayscale().to_rgba8()
}

/// 颜色反转
fn process_invert(img: &DynamicImage) -> RgbaImage {
    let mut img_rgba = img.to_rgba8();
    imageops::invert(&mut img_rgba);
    img_rgba
}

/// 高斯模糊
fn process_blur(img: &DynamicImage) -> RgbaImage {
    img.blur(5.0).to_rgba8()
}

/// 旋转90度
fn process_rotate90(img: &DynamicImage) -> RgbaImage {
    img.rotate90().to_rgba8()
}

/// Sobel边缘检测
fn process_canny_edge(img: &DynamicImage) -> RgbaImage {
    let gray_img = img.to_luma8();
    let mut edges = gray_img.clone();
    
    for y in 1..gray_img.height()-1 {
        for x in 1..gray_img.width()-1 {
            let mut gx = 0i16;
            let mut gy = 0i16;
            
            for dy in -1..=1 {
                for dx in -1..=1 {
                    let pixel = gray_img.get_pixel(
                        (x as i32 + dx) as u32, 
                        (y as i32 + dy) as u32
                    )[0] as i16;
                    
                    match dx {
                        -1 => gx -= pixel,
                        1 => gx += pixel,
                        _ => {}
                    }
                    
                    match dy {
                        -1 => gy -= pixel,
                        1 => gy += pixel,
                        _ => {}
                    }
                }
            }
            
            let magnitude = (gx.abs() + gy.abs()).min(255) as u8;
            edges.put_pixel(x, y, image::Luma([magnitude]));
        }
    }
    
    DynamicImage::ImageLuma8(edges).to_rgba8()
}

/// 处理图像字节数据 (Python 接口)
#[pyfunction]
fn process_image_bytes(py: Python, input: &[u8], command: &str) -> PyResult<Py<PyBytes>> {
    let img = image::load_from_memory(input)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
    
    let result = match command {
        "grayscale" => process_grayscale(&img),
        "invert" => process_invert(&img),
        "blur" => process_blur(&img),
        "rotate90" => process_rotate90(&img),
        "canny_edge" => process_canny_edge(&img),
        _ => return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("Unknown command: {}", command)
        )),
    };

    // 编码为PNG
    let mut buffer = Vec::new();
    result.write_to(&mut std::io::Cursor::new(&mut buffer), image::ImageOutputFormat::Png)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
    
    Ok(PyBytes::new(py, &buffer).into())
}

/// 处理 numpy 数组图像 (Python 接口)
#[pyfunction]
fn process_numpy_image(py: Python, image: &PyArrayDyn<u8>, command: &str) -> PyResult<Py<PyArrayDyn<u8>>> {
    // 获取 numpy 数组的形状和数据
    let shape = image.shape();
    if shape.len() != 3 || (shape[2] != 3 && shape[2] != 4) {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "Invalid image shape. Expected (height, width, 3) or (height, width, 4)"
        ));
    }
    
    let height = shape[0] as u32;
    let width = shape[1] as u32;
    let channels = shape[2];
    
    // 创建动态图像
    let img_data = image.as_slice().map_err(|e| 
        PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
    )?;
    
    let img = match channels {
        3 => DynamicImage::ImageRgb8(
            image::ImageBuffer::from_raw(width, height, img_data.to_vec())
                .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    "Failed to create RGB image buffer"))?
        ),
        4 => DynamicImage::ImageRgba8(
            image::ImageBuffer::from_raw(width, height, img_data.to_vec())
                .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    "Failed to create RGBA image buffer"))?
        ),
        _ => unreachable!()
    };
    
    // 处理图像
    let result = match command {
        "grayscale" => process_grayscale(&img),
        "invert" => process_invert(&img),
        "blur" => process_blur(&img),
        "rotate90" => process_rotate90(&img),
        "canny_edge" => process_canny_edge(&img),
        _ => return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            format!("Unknown command: {}", command)
        )),
    };
    
    // 获取结果尺寸
    let (result_width, result_height) = result.dimensions();
    let result_channels = 4; // 总是返回RGBA
    
    // 创建新的numpy数组
    let result_array = PyArray::from_vec(
        py,
        result.into_raw(),
        [result_height as usize, result_width as usize, result_channels]
    ).into_dyn();
    
    Ok(result_array.into_py(py))
}

/// Python 模块定义
#[pymodule]
fn image_processing(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(process_image_bytes, m)?)?;
    m.add_function(wrap_pyfunction!(process_numpy_image, m)?)?;
    Ok(())
}